{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e8b27860",
   "metadata": {},
   "source": [
    "# What\n",
    "分类任务，支持两种模式\n",
    "1. Folder模式，需要输入`train`, `valid`两个测试集对应的目录。`labels.txt`，需要训练的label，里面每个类别一行。\n",
    "2. List模式，需要输入`train`, `valid`两个测试集对应的训练文件，每行一个样本。`labels.txt`是可选参数，里面每个类别一行。`data_pattern`一个通用的目录，与train、val中的第一列进行拼接。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f33efa53",
   "metadata": {},
   "source": [
    "### 获取训练数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3913fb27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "n_classes = 2\n",
    "mydir = r'C:\\Users\\onekey\\Documents\\WeChat Files\\wxid_v5pywcdqb00a12\\FileStorage\\File\\2022-02\\T2ROI'\n",
    "img_list = [(i, f\"{i.replace('.nii', '_Merge.nii')}\") for i in os.listdir(mydir) if 'Merge' not in i]\n",
    "train_img_list = img_list[:30]\n",
    "train_img_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7df8892",
   "metadata": {},
   "outputs": [],
   "source": [
    "from onekey_algo.custom.components.Radiology import get_image_mask_from_dir, diagnose_3d_image_mask_settings\n",
    "\n",
    "n_classes = 2\n",
    "my_dir = r'C:\\Users\\onekey\\Documents\\Task06_Lung'\n",
    "images, masks = get_image_mask_from_dir(my_dir, images='imagesTr', masks='labelsTr')\n",
    "\n",
    "# diagnose_3d_image_mask_settings(images, masks)\n",
    "print(f'获取到{len(images)}个样本。')\n",
    "train_img_list = list(zip(images, masks))\n",
    "mydir = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e4ab32d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "\n",
    "n_classes = 2 \n",
    "mydir = r'C:\\Users\\onekey\\Project\\ROI1'\n",
    "\n",
    "img_list = []\n",
    "for i in os.listdir(mydir):\n",
    "    if os.path.exists(os.path.join(mydir, i, 'data.nii.gz')):\n",
    "        if os.path.exists(os.path.join(mydir, i, 'mask.nii.gz')):\n",
    "            img_list.append([os.path.join(mydir, i, 'data.nii.gz'), os.path.join(mydir, i, 'mask.nii.gz')])\n",
    "        elif os.path.exists(os.path.join(mydir, i, 'mask.nrrd')):\n",
    "            img_list.append([os.path.join(mydir, i, 'data.nii.gz'), os.path.join(mydir, i, 'mask.nrrd')])\n",
    "mydir = ''\n",
    "print(f'获取到{len(img_list)}个样本。')\n",
    "random.shuffle(img_list)\n",
    "train_img_list = img_list[:int(len(img_list) * 0.7)]\n",
    "test_img_list = img_list[int(len(img_list) * 0.7):]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0873dca7",
   "metadata": {},
   "source": [
    "### 训练模式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c449c95c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from onekey_algo.segmentation3D.transfer_learning.run_3dsegmentation import main as which3D_main\n",
    "from collections import namedtuple\n",
    "\n",
    "# 设置参数\n",
    "class params:\n",
    "    data_root = mydir\n",
    "    img_list = train_img_list\n",
    "    n_seg_classes = n_classes\n",
    "    learning_rate = 0.1\n",
    "    num_workers = 6\n",
    "    batch_size = 2\n",
    "    phase = 'train'\n",
    "    save_intervals = 10\n",
    "    n_epochs = 300\n",
    "    input_D = 28\n",
    "    input_H = 448\n",
    "    input_W = 448\n",
    "    resume_path = ''\n",
    "    pretrain_path = r'C:/Users/onekey/Documents/pretrain/resnet_50_23dataset.pth'\n",
    "    new_layer_names = [\n",
    "        'conv_seg']  # default=['upsample1', 'cmp_layer3', 'upsample2', 'cmp_layer2', 'upsample3', 'cmp_layer1', 'upsample4', 'cmp_conv1', 'conv_seg']\n",
    "    no_cuda = False\n",
    "    gpu_id = [0]\n",
    "    model = 'resnet'\n",
    "    model_depth = 50\n",
    "    resnet_shortcut = 'B'\n",
    "    manual_seed = 1\n",
    "    ci_test = False\n",
    "    attr = {}\n",
    "\n",
    "    def __setattr__(self, key, value):\n",
    "        self.attr[key] = value\n",
    "# 训练模型\n",
    "params.save_folder = \"./models/{}_{}\".format(params.model, params.model_depth)\n",
    "which3D_main(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfbb2e69",
   "metadata": {},
   "source": [
    "### 测试模式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f51e166",
   "metadata": {},
   "outputs": [],
   "source": [
    "from onekey_algo.segmentation3D.transfer_learning.eval_3dsegmentation import main as which3D_main\n",
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='0'\n",
    "\n",
    "test_img_list = train_img_list[:5]\n",
    "# 设置参数\n",
    "class params:\n",
    "    data_root = mydir\n",
    "    img_list = test_img_list\n",
    "    n_seg_classes = n_classes\n",
    "    learning_rate = 0.001\n",
    "    num_workers = 1\n",
    "    batch_size = 1\n",
    "    phase = 'test'\n",
    "    save_intervals = 10\n",
    "    n_epochs = 30\n",
    "    input_D = 28\n",
    "    input_H = 448\n",
    "    input_W = 448\n",
    "    resume_path = './models/resnet_50_epoch_14_batch_225.pth'\n",
    "    pretrain_path = r'C:\\Users\\onekey\\Documents\\pretrain\\resnet_50_23dataset.pth'\n",
    "    new_layer_names = [\n",
    "        'conv_seg']  # default=['upsample1', 'cmp_layer3', 'upsample2', 'cmp_layer2', 'upsample3', 'cmp_layer1', 'upsample4', 'cmp_conv1', 'conv_seg']\n",
    "    no_cuda = False\n",
    "    gpu_id = [0]\n",
    "    model = 'resnet'\n",
    "    model_depth = 50\n",
    "    resnet_shortcut = 'B'\n",
    "    manual_seed = 1\n",
    "    ci_test = False\n",
    "    attr = {}\n",
    "\n",
    "    def __setattr__(self, key, value):\n",
    "        self.attr[key] = value\n",
    "# 训练模型\n",
    "params.save_folder = \"./models/{}_{}\".format(params.model, params.model_depth)\n",
    "which3D_main(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "771684b1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
